import cv2
import numpy as np
from matplotlib import pyplot as plt

# 计算距离
def cal_distance(pa, pb):
    from math import sqrt
    dis = sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)
    return dis

# 返回高斯低通滤波矩阵
def get_gauss_lowpass(shape, d0):
    mat = np.zeros(shape)
    center_point = (shape[0]//2, shape[1]//2)
    for i in range(shape[0]):
        for j in range(shape[1]):
            dis = cal_distance(center_point, (i, j))
            mat[i, j] = np.exp(-(dis ** 2) / (2 * (d0 ** 2)))

    return mat

# 返回高斯高通滤波矩阵
def get_gauss_highpass(shape, d0):
    return 1 - get_gauss_lowpass(shape, d0)


image = cv2.imread('beauty.jpg', cv2.IMREAD_GRAYSCALE)
# 变换到频率域
f_image = np.fft.fft2(image)
f_image = np.fft.fftshift(f_image)
spectrum = 20 * np.log(np.abs(f_image))
# 频率域滤波
filter_mat = get_gauss_lowpass(image.shape, 50)
f_filter_image = f_image * filter_mat
filter_spectrum = 20 * np.log(np.abs(f_filter_image))
# 反变换回空间域
f_filter_image = np.fft.ifftshift(f_filter_image)
filter_image = np.fft.ifft2(f_filter_image)
filter_image = np.abs(filter_image)
# 显示结果图像
plt.subplot(221), plt.imshow(image, cmap="gray")
plt.title("Src image"), plt.xticks([]), plt.yticks([])
plt.subplot(222), plt.imshow(spectrum, cmap="bwr")
plt.title('Src spectrum'), plt.xticks([]), plt.yticks([])
plt.subplot(223), plt.imshow(filter_image, cmap='gray')
plt.title("Filter image"), plt.xticks([]), plt.yticks([])
plt.subplot(224), plt.imshow(filter_spectrum, cmap="bwr")
plt.title('Filter spectrum'), plt.xticks([]), plt.yticks([])
plt.show()
